# -- encoding:utf-8 --

import tensorflow as tf

if __name__ == '__main__':
    network_name = "TEXT_CNN"

    # 修改路径即可
    ckpt_path = r"..\wenshu\model\text_cnn\model.ckpt-230"
    pb_path = "./wenshu.model.pb"

    # 下面三个为网络结构中输入以及输出Tensor的名称字符串
    input_text_tensor_name = "{}/placeholders/input_word_id".format(network_name)
    output_predict_tensor_name = "{}/project/predictions".format(network_name)
    output_predict_probability_tensor_name = "{}/project/probability".format(network_name)
    # 在构建pb文件的时候，输出的tensor字符串名称列表(字符串后面不能加:0这种后缀)
    output_node_names = [
        input_text_tensor_name,
        output_predict_tensor_name,
        output_predict_probability_tensor_name
    ]

    # 一、模型构建、训练(模型执行图及模型参数恢复)
    # 0. 文件检查
    if not tf.gfile.Exists('{}.meta'.format(ckpt_path)):
        raise Exception("meta文件不存在，请检查输入路径:{}".format(ckpt_path))
    if not tf.gfile.Exists('{}.index'.format(ckpt_path)):
        raise Exception("index文件不存在，请检查输入路径:{}".format(ckpt_path))

    # 1. 构建会话
    graph = tf.get_default_graph()
    sess = tf.Session()
    # 2. 执行图恢复
    saver = tf.train.import_meta_graph('{}.meta'.format(ckpt_path))
    # 3. 模型参数恢复
    saver.restore(sess, ckpt_path)
    # 4. 获取关注的Tensor对象(也就是网络的输入、输出)
    inputs = graph.get_tensor_by_name("{}:0".format(input_text_tensor_name))
    predictions = graph.get_tensor_by_name("{}:0".format(output_predict_tensor_name))
    probability = graph.get_tensor_by_name("{}:0".format(output_predict_probability_tensor_name))
    print(inputs, predictions, probability)

    # 二、将模型输出为pb文件(仅保存output_node_names涉及到的执行图)
    convert_graph_def = tf.graph_util.convert_variables_to_constants(sess=sess,
                                                                     input_graph_def=sess.graph.as_graph_def(),
                                                                     output_node_names=output_node_names)
    # 将对象序列化为二进制字符串输出即可
    with tf.gfile.GFile(pb_path, 'wb') as writer:
        writer.write(convert_graph_def.SerializeToString())
